import cv2
import numpy as np
import mss
import pyautogui
import time
import win32gui
from ultralytics import YOLO
import random
import os
from datetime import datetime


class AimbotModule:
    def __init__(self, tasks):
        model_path = r'H:\yol\yolo\v11\runs\detect\train18\weights\best.pt'
        self.model = YOLO(model_path).to('cpu')  # 强制 CPU 模式
        self.screenshot_dir = r'H:\yol\ultralytics\datasets\mycp\images'  # 截图保存目录
        os.makedirs(self.screenshot_dir, exist_ok=True)  # 创建目录（如果不存在）

        # 自动查找梦幻西游窗口
        self.target_hwnd = self._find_game_window()
        if self.target_hwnd:
            window_title = win32gui.GetWindowText(self.target_hwnd)
            print(f"✅ 成功绑定窗口：{window_title}")
        else:
            raise ValueError("❌ 未找到标题包含'梦幻西游'的窗口")

        self.tasks = tasks
        self._update_window_info()
        self.work_area = (
            self.window_left,
            self.window_top,
            self.window_left + self.window_width,
            self.window_top + self.window_height
        )
        self.monitor = {
            "top": self.window_top,
            "left": self.window_left,
            "width": self.window_width,
            "height": self.window_height
        }
        self.window_x = self.window_left + self.window_width  # 监视窗口在游戏窗口右侧
        self.window_y = 0  # 监视窗口顶部对齐
        self.init_monitor_window()  # 初始化监视窗口
        self.mouse_last_in_area_time = time.time()  # 初始化鼠标最后在区域内的时间
        self.last_known_target = None  # 记录最后一次已知的目标位置
        self.clicked = False  # 标记是否已点击
        self.clicked_class = None  # 记录被点击的类别
        self.last_success_time = time.time()  # 记录上次任务成功的时间
        self.locked_target = None  # 锁定的目标
        self.target_lost_time = None  # 目标丢失的时间
        self.stored_target = None  # 存储目标坐标
        self.last_crosshair_pos = None  # 记录上一次准星的位置
        self.crosshair_still_time = 0  # 准星静止的时间

    def _find_game_window(self):
        """遍历所有窗口，查找标题包含'梦幻西游'的窗口"""
        target_hwnds = []

        def enum_callback(hwnd, extra):
            title = win32gui.GetWindowText(hwnd)
            if "梦幻西游 ONLINE" in title:
                extra.append(hwnd)
                print(f"🔍 发现匹配窗口：{title}")  # 打印所有匹配窗口
            return True

        win32gui.EnumWindows(enum_callback, target_hwnds)
        return target_hwnds[0] if target_hwnds else None

    def _update_window_info(self):
        """实时更新窗口坐标和尺寸"""
        if self.target_hwnd:
            left, top, right, bottom = win32gui.GetWindowRect(self.target_hwnd)
            self.window_left, self.window_top = left, top
            self.window_width, self.window_height = right - left, bottom - top

    def init_monitor_window(self):
        """初始化监视窗口（保留原有窗口设置）"""
        cv2.namedWindow('Aimbot Monitor', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Aimbot Monitor', 512, 400)
        cv2.moveWindow('Aimbot Monitor', self.window_x, self.window_y)

    def start(self, distance_threshold=3, delay_between_tasks=1):
        """启动瞄准系统（完整保留监视窗口绘制）"""
        self.stop_flag = False
        sct = mss.mss()
        target_reacquire_time = 2  # 目标重新获取的时间阈值
        crosshair_still_threshold = 0.5  # 准星静止的时间阈值

        for task in self.tasks:
            target_class = task['target_class']
            success_class = task['success_class']
            refresh_keys = task.get('refresh_keys', [])
            click_offset = task.get('click_offset', (0, 0))
            click_type_num = task.get('click_type', 1)  # 默认左键点击

            if click_type_num == 1:
                click_type = 'left'
            elif click_type_num == 2:
                click_type = 'right'
            else:
                click_type = None

            self.target_class = target_class
            self.success_class = success_class
            print(f"开始点击目标 id {target_class}，成功 ID 为 {success_class}，点击类型：{click_type}")

            if refresh_keys:
                print(f"按下组合键 {refresh_keys} 刷新场景")
                pyautogui.hotkey(*refresh_keys)
                time.sleep(1)  # 等待刷新完成

            while not self.stop_flag:
                self._update_window_info()
                current_time = time.time()
                mouse_x, mouse_y = pyautogui.position()

                # 截图与检测（提前检测类别0是否存在）
                sct_img = sct.grab(self.monitor)
                img_np = np.array(sct_img)
                # 将4通道图像转换为3通道图像
                img_np = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
                results = self.model.predict(img_np, verbose=False)
                detected_classes = {int(box.cls[0]) for result in results for box in result.boxes.numpy()}

                # 类别0消失处理：小幅度晃动鼠标
                if 0 not in detected_classes:
                    self.save_screenshot(img_np, "class0_disappear")
                    self.shake_mouse()  # 触发鼠标晃动
                    time.sleep(0.5)  # 等待画面更新

                # 鼠标位置处理（原有逻辑）
                if not (self.work_area[0] <= mouse_x <= self.work_area[2] and
                        self.work_area[1] <= mouse_y <= self.work_area[3]):
                    if current_time - self.mouse_last_in_area_time >= 1:
                        center_x = self.work_area[0] + self.window_width // 2
                        center_y = self.work_area[1] + self.window_height // 2
                        # 添加随机量
                        random_x = random.randint(-10, 10)
                        random_y = random.randint(-10, 10)
                        pyautogui.moveTo(center_x + random_x, center_y + random_y, duration=0.1)
                else:
                    self.mouse_last_in_area_time = current_time

                # 检查是否连续 5 秒任务未成功
                if current_time - self.last_success_time >= 5:
                    center_x = self.work_area[0] + self.window_width // 2
                    center_y = self.work_area[1] + self.window_height // 2
                    # 添加随机量
                    random_x = random.randint(-10, 10)
                    random_y = random.randint(-10, 10)
                    pyautogui.moveTo(center_x + random_x, center_y + random_y, duration=0.1)
                    self.clicked = False
                    self.clicked_class = None
                    self.locked_target = None
                    self.target_lost_time = None
                    self.last_success_time = current_time
                    self.stored_target = None
                    print("❌ 连续 5 秒任务未成功，重置鼠标位置并重启任务。")

                try:
                    img = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
                    results = self.model.predict(img, verbose=False)

                    # 收集并输出类信息（新增功能）
                    detected_classes = {int(box.cls[0]) for result in results for box in result.boxes.numpy()}
                    if detected_classes:
                        class_info = []
                        for cls_id in detected_classes:
                            class_name = self.model.names.get(cls_id, f"未知类别(ID:{cls_id})")
                            class_info.append(f"{class_name}({cls_id})")
                        print(f"📌 检测到的类：{', '.join(class_info)}")

                    # 目标检测与绘制（原有逻辑，保留检测框绘制）
                    new_enemy_pos = []
                    new_crosshair = None
                    for result in results:
                        for box in result.boxes.numpy():
                            x1, y1, x2, y2 = box.xyxy[0].astype(int)
                            cls_id = int(box.cls[0])

                            # 十字准星与目标位置（原有逻辑）
                            if cls_id == 0:
                                new_crosshair = (x1, y1)
                            elif cls_id == self.target_class:
                                new_enemy_pos.append(((x1 + x2) // 2, (y1 + y2) // 2))

                            # 绘制检测框和类 ID（原有功能保留）
                            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            cv2.putText(img, f"{cls_id}", (x1, y1 - 10),
                                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

                    # 瞄准逻辑（修改部分）
                    self.enemy_positions = new_enemy_pos
                    self.crosshair_pos = new_crosshair
                    if self.crosshair_pos and (self.enemy_positions or self.last_known_target):
                        if self.locked_target and self.locked_target in self.enemy_positions:
                            closest = self.locked_target
                        elif self.enemy_positions:
                            closest = min(self.enemy_positions, key=lambda p: np.linalg.norm(
                                np.array(p) - np.array(self.crosshair_pos)))
                            self.locked_target = closest
                            self.target_lost_time = None
                        elif self.locked_target and self.target_lost_time and (
                                current_time - self.target_lost_time) < target_reacquire_time:
                            closest = self.locked_target
                        else:
                            closest = self.last_known_target
                            self.locked_target = None
                            self.target_lost_time = None

                        if closest not in self.enemy_positions and self.locked_target:
                            if not self.target_lost_time:
                                self.target_lost_time = current_time
                        elif closest in self.enemy_positions and self.target_lost_time:
                            self.target_lost_time = None

                        if closest:
                            m, n = closest
                            self.stored_target = (m, n)
                        elif self.stored_target:
                            m, n = self.stored_target
                        else:
                            continue

                        a, b = click_offset
                        x, y = self.window_left + self.crosshair_pos[0], self.window_top + self.crosshair_pos[1]
                        actual_m = self.window_left + m + a
                        actual_n = self.window_top + n + b

                        # 鼠标相对移动 (m + a - x, n + b - y) * 0.8
                        dx = int((actual_m - x) * 0.8)
                        dy = int((actual_n - y) * 0.8)
                        pyautogui.moveRel(dx, dy, duration=0.01)

                        # 计算类别 0 和目标类别的距离
                        distance = np.linalg.norm(np.array([x, y]) - np.array([actual_m, actual_n]))

                        # 检查准星是否静止
                        if self.last_crosshair_pos == self.crosshair_pos:
                            self.crosshair_still_time += 0.01
                            if self.crosshair_still_time >= crosshair_still_threshold:
                                if click_type:
                                    pyautogui.click(button=click_type)
                                    self.clicked = True
                                    self.clicked_class = self.target_class
                                    self.stored_target = None  # 点击瞬间清空存储的坐标
                                    print(
                                        f"已经对目标 id {self.target_class} 执行{click_type}点击，正在检测此目标是否还在。")
                                time.sleep(1)  # 延迟 1 秒

                                # 重新进行检测
                                sct_img = sct.grab(self.monitor)
                                img_np = np.array(sct_img)
                                img_np = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
                                results = self.model.predict(img_np, verbose=False)
                                detected_classes = {int(box.cls[0]) for result in results for box in
                                                    result.boxes.numpy()}
                                if detected_classes:
                                    class_info = []
                                    for cls_id in detected_classes:
                                        class_name = self.model.names.get(cls_id, f"未知类别(ID:{cls_id})")
                                        class_info.append(f"{class_name}({cls_id})")
                                    print(f"📌 点击后重新检测到的类：{', '.join(class_info)}")
                                if self.success_class is not None and self.success_class in detected_classes:
                                    print(
                                        f"✅ 目标 id {self.target_class} 点击成功，检测到成功 ID {self.success_class}，结束此任务。")
                                    self.last_success_time = current_time
                                    self.clicked = False
                                    self.clicked_class = None
                                    self.locked_target = None
                                    self.target_lost_time = None
                                    self.stored_target = None
                                    self.last_crosshair_pos = None
                                    self.crosshair_still_time = 0
                                    time.sleep(delay_between_tasks)
                                    break
                                elif self.target_class not in detected_classes:
                                    print(f"✅ 目标 id {self.target_class} 已消失，点击成功，结束此任务。")
                                    self.last_success_time = current_time
                                    self.clicked = False
                                    self.clicked_class = None
                                    self.locked_target = None
                                    self.target_lost_time = None
                                    self.stored_target = None
                                    self.last_crosshair_pos = None
                                    self.crosshair_still_time = 0
                                    time.sleep(delay_between_tasks)
                                    break
                        else:
                            self.last_crosshair_pos = self.crosshair_pos
                            self.crosshair_still_time = 0

                        # 新增逻辑：即使目标未被检测到，准星到了大致位置也尝试点击
                        if self.locked_target and distance < distance_threshold:
                            if click_type:
                                pyautogui.click(button=click_type)
                                self.clicked = True
                                self.clicked_class = self.target_class
                                self.stored_target = None  # 点击瞬间清空存储的坐标
                                print(f"已经对目标 id {self.target_class} 执行{click_type}点击，正在检测此目标是否还在。")
                            time.sleep(1)  # 延迟 1 秒

                            # 重新进行检测
                            sct_img = sct.grab(self.monitor)
                            img_np = np.array(sct_img)
                            img_np = cv2.cvtColor(img_np, cv2.COLOR_BGRA2BGR)
                            results = self.model.predict(img_np, verbose=False)
                            detected_classes = {int(box.cls[0]) for result in results for box in result.boxes.numpy()}
                            if detected_classes:
                                class_info = []
                                for cls_id in detected_classes:
                                    class_name = self.model.names.get(cls_id, f"未知类别(ID:{cls_id})")
                                    class_info.append(f"{class_name}({cls_id})")
                                print(f"📌 点击后重新检测到的类：{', '.join(class_info)}")
                            if self.success_class is not None and self.success_class in detected_classes:
                                print(
                                    f"✅ 目标 id {self.target_class} 点击成功，检测到成功 ID {self.success_class}，结束此任务。")
                                self.last_success_time = current_time
                                self.clicked = False
                                self.clicked_class = None
                                self.locked_target = None
                                self.target_lost_time = None
                                self.stored_target = None
                                self.last_crosshair_pos = None
                                self.crosshair_still_time = 0
                                time.sleep(delay_between_tasks)
                                break
                            elif self.target_class not in detected_classes:
                                print(f"✅ 目标 id {self.target_class} 已消失，点击成功，结束此任务。")
                                self.last_success_time = current_time
                                self.clicked = False
                                self.clicked_class = None
                                self.locked_target = None
                                self.target_lost_time = None
                                self.stored_target = None
                                self.last_crosshair_pos = None
                                self.crosshair_still_time = 0
                                time.sleep(delay_between_tasks)
                                break

                        if self.clicked and self.success_class is not None:
                            if self.success_class in detected_classes:
                                print(
                                    f"✅ 目标 id {self.target_class} 点击成功，检测到成功 ID {self.success_class}，结束此任务。")
                                self.last_success_time = current_time
                                self.clicked = False
                                self.clicked_class = None
                                self.locked_target = None
                                self.target_lost_time = None
                                self.stored_target = None
                                self.last_crosshair_pos = None
                                self.crosshair_still_time = 0
                                time.sleep(delay_between_tasks)
                                break

                    # 显示监视窗口（原有功能，确保窗口绘制）
                    cv2.imshow('Aimbot Monitor', img)
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        self.stop()

                except Exception as e:
                    print(f"❌ 错误: {e}")
                    self.save_screenshot(img, "error_occurred")  # 错误时也截图

                time.sleep(0.01)  # 控制循环频率

        sct.close()
        cv2.destroyAllWindows()

    def save_screenshot(self, img, reason):
        """保存截图到指定目录"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{reason}_{timestamp}.png"
        path = os.path.join(self.screenshot_dir, filename)
        cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_BGRA2BGR))
        print(f"📸 截图保存：{path}（原因：{reason}）")

    def shake_mouse(self):
        """小幅度晃动鼠标（先右后左，各5像素）"""
        current_x, current_y = pyautogui.position()
        # 向右晃动
        pyautogui.moveRel(5, 0, duration=0.1)
        time.sleep(0.1)
        # 向左晃动
        pyautogui.moveRel(-10, 0, duration=0.1)
        time.sleep(0.1)
        # 恢复原位
        pyautogui.moveRel(5, 0, duration=0.1)
        print("鼠标小幅度晃动，尝试重现类别0")

    def stop(self):
        self.stop_flag = True
        cv2.destroyWindow('Aimbot Monitor')  # 关闭监视窗口
